import asyncio
import asyncpg
from urllib.parse import quote_plus
from src.core.database import db_manager
from src.core.config import settings
from src.core.logger import get_logger

logger = get_logger(__name__)

async def setup_database_and_user():
    """使用管理员凭据创建数据库和用户，并授予权限"""
    encoded_admin_password = quote_plus(settings.admin_db_password)
    admin_db_url = (
        f"postgresql://{settings.admin_db_user}"
        f":{encoded_admin_password}@{settings.database_host}"
        f":{settings.database_port}/postgres" # Connect to default 'postgres' db for admin ops
    )
    
    conn = None
    try:
        logger.info(f"尝试以管理员身份连接到数据库: {settings.database_host}:{settings.database_port}/postgres")
        conn = await asyncpg.connect(admin_db_url)
        logger.info("管理员连接成功")

        # 检查并创建数据库
        db_name = settings.database_name
        check_db_exists = await conn.fetchval(f"SELECT 1 FROM pg_database WHERE datname = '{db_name}'")
        if not check_db_exists:
            logger.info(f"数据库 '{db_name}' 不存在，正在创建...")
            await conn.execute(f"CREATE DATABASE {db_name}")
            logger.info(f"数据库 '{db_name}' 创建成功")
        else:
            logger.info(f"数据库 '{db_name}' 已存在")

        # 检查并创建用户
        db_user = settings.database_user
        db_password = settings.database_password
        check_user_exists = await conn.fetchval(f"SELECT 1 FROM pg_roles WHERE rolname = '{db_user}'")
        if not check_user_exists:
            logger.info(f"用户 '{db_user}' 不存在，正在创建...")
            await conn.execute(f"CREATE USER {db_user} WITH PASSWORD '{db_password}'")
            logger.info(f"用户 '{db_user}' 创建成功")
        else:
            logger.info(f"用户 '{db_user}' 已存在")

        # 授予用户对数据库的权限
        logger.info(f"授予用户 '{db_user}' 对数据库 '{db_name}' 的所有权限...")
        await conn.execute(f"GRANT ALL PRIVILEGES ON DATABASE {db_name} TO {db_user}")
        logger.info("权限授予成功")

    except asyncpg.exceptions.DuplicateDatabaseError:
        logger.info(f"数据库 '{settings.database_name}' 已存在，跳过创建")
    except asyncpg.exceptions.DuplicateObjectError:
        logger.info(f"用户 '{settings.database_user}' 已存在，跳过创建")
    except Exception as e:
        logger.error(f"数据库或用户设置失败: {e}", error=str(e))
        raise
    finally:
        if conn:
            await conn.close()

async def main():
    await setup_database_and_user()
    await db_manager.initialize()
    await db_manager.create_tables()
    await db_manager.close()

if __name__ == "__main__":
    asyncio.run(main())